import json
import os
import time
from tenacity import retry, stop_after_attempt, wait_random_exponential

from tool import get_config, api_request

CONFIG = get_config('bianxie')

# 创建保存目录
os.makedirs(CONFIG['save_dir'], exist_ok=True)

# 带重试机制的请求装饰器
@retry(stop=stop_after_attempt(CONFIG['retries']), 
       wait=wait_random_exponential(multiplier=1, max=CONFIG['max_delay']))
def api_request_v1(model, prompt, attempt):
    """执行API请求"""
    api_request(model, prompt, attempt, CONFIG)

def process_prompts():
    # 加载提示词
    with open('data/prompt/single_prompts.json', 'r', encoding='utf-8') as f:
        prompts = json.load(f)

    # 遍历所有问题
    for prompt_key, prompt_data in prompts.items():
        question = prompt_data
        
        # 遍历所有模型
        for model in CONFIG['models']:
            # 重复50次请求
            for i in range(1, 51):
                model_name = model.split("/")[-1] if "/" in model else model
                filename = f"{CONFIG['save_dir']}/{prompt_key}_{model_name}_{i}.json"
                
                # 跳过已完成的请求
                if os.path.exists(filename):
                    continue
                
                # 执行请求
                try:
                    start_time = time.time()
                    answer = api_request_v1(model, question, i)
                    
                    # 保存结果
                    with open(filename, 'w', encoding='utf-8') as f:
                        json.dump({
                            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                            "model": model,
                            "prompt": question,
                            "response": answer,
                            "latency": round(time.time() - start_time, 2)
                        }, f, ensure_ascii=False, indent=2)
                    
                    print(f"Success: {filename}")
                    
                except Exception as e:
                    print(f"Final failure: {filename} - {str(e)}")
                    # 保存错误信息
                    with open(filename, 'w', encoding='utf-8') as f:
                        json.dump({
                            "error": str(e),
                            "attempt": i
                        }, f, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    process_prompts()